//
// Created by kai-xu on 2023/2/28.
//

#ifndef PRO_MAT_INCLUDE_MATRIXVT_H_
#define PRO_MAT_INCLUDE_MATRIXVT_H_

#include <Eigen/Core>
#include <Eigen/Geometry>
#include <iostream>
#include <vector>

#include "common/basic_type.h"
#include "common/file_helper.h"
#include "common/utility.h"

namespace autodrive {
namespace perception {
struct CustomCoorPoint {
  /// @brief static_cast<int32_t>(x-coordinate).
  int32_t x;
  /// @brief static_cast<int32_t>(y-coordinate).
  int32_t y;
  /// @brief static_cast<int32_t>(z-coordinate).
  int32_t z;
  /// @brief x-coordinate.
  float xf;
  /// @brief y-coordinate.
  float yf;
  /// @brief  z-coordinate.
  float zf;
};

struct CamCrop {
  /// @brief x-coordinate of upper left corner (lowest memory address).
  float x;
  /// @brief y-coordinate of upper left corner (lowest memory address).
  float y;
  /// @brief  x-coordinate of downer right corner (lowest memory address).
  float x1;
  /// @brief y-coordinate of downer right corner (lowest memory address).
  float y1;
};

/// @brief matrixvt parameters
struct MatrixVTParameters {
  /// @brief batch size  feature.shape[0].  BNDHWC
  int32_t batch_size;
  /// @brief cam number  feature.shape[1].
  int32_t n_cam;
  /// @brief feature.shape[2].
  int32_t dz;
  /// @brief feature.shape[3].
  int32_t height;
  /// @brief feature.shape[4]
  int32_t width;
  /// @brief bev feature width.
  int32_t out_width;
  /// @brief bev feature height.
  int32_t out_height;
  /// @brief strid x.
  int32_t sd_x;
  /// @brief strid y.
  int32_t sd_y;
  /// @brief data_scale.
  int16_t scale;
  /// @brief bev voxel size.
  std::vector<float> voxel_size;
  /// @brief bev range.
  std::vector<float> net_range;
  /// @brief crop paramter for aug_mat.
  std::vector<CamCrop> rect_list;
  /// @brief resize paramter for aug_mat
  std::vector<float> resize_list;
};

bool gen_aug_ego_pose(std::vector<double> &augMat, std::vector<double> &egoMat,
                      std::vector<double> const &Cam_Intrins,
                      std::vector<double> const &Cam_Extrins,
                      std::vector<float> const &resize_scale,
                      std::vector<CamCrop> const &Cam_Crop);

int32_t cal_coor_cpu(CustomCoorPoint *key_value, double *aug_mat,
                     double *ego_mat, std::vector<float> const &voxel_size,
                     std::vector<float> const &coors_range, int32_t sd_x,
                     int32_t sd_y, int32_t n_cam, int32_t dz, int32_t width,
                     int32_t height);

int32_t coor_reduce_hdim_cpu(CustomCoorPoint *o_coor_value,
                             const CustomCoorPoint *i_coor_value, int32_t bs,
                             int32_t nc, int32_t dz, int32_t height,
                             int32_t width);

int32_t coor_permute_cpu(CustomCoorPoint *o_coor_value,
                         const CustomCoorPoint *i_coor_value, int32_t *idx_ptr,
                         int32_t bs, int32_t nc, int32_t dz, int32_t width,
                         int32_t out_height, int32_t out_width,
                         int32_t invalid_x, int32_t invalid_y);

template <typename T>
int32_t coor_map_cpu(T *circle_map, T *ray_map, const int32_t *idx_ptr,
                     int32_t bs, int32_t dz, int32_t nc, int32_t width,
                     int32_t out_height, int32_t out_width, int16_t scale,
                     int32_t invalid_x, int32_t invalid_y) {
  constexpr float quantifyOffset{0.5};
  // get circle map
  for (int32_t k{0}; k < bs; ++k) {
    for (int32_t j{0}; j < dz; ++j) {
      for (int32_t n{0}; n < nc * width; ++n) {
        auto adr = k * dz * nc * width + n * dz + j;
        circle_map[k * dz * out_height * out_width +
                   j * out_height * out_width + idx_ptr[adr]] =
            quantify<float, T>(1.0, scale, quantifyOffset);
      }
    }
  }

  // get ray map
  for (int32_t i{0}; i < bs; ++i) {
    for (int32_t j{0}; j < nc * width; ++j) {
      auto adr = i * nc * width * dz + j * dz;
      for (int32_t n{0}; n < dz; ++n) {
        ray_map[i * nc * width * out_height * out_width +
                j * out_height * out_width + idx_ptr[adr + n]] =
            quantify<float, T>(1.0, scale, quantifyOffset);
        ;
      }
    }
  }

  // clean invalid value
  const int invalid_idx = invalid_y * out_width + invalid_x;
  const int ray_map_idx = bs * nc * width;
  for (int32_t idx{0}; idx < ray_map_idx; ++idx) {
    int cw = idx % (nc * width);
    int b = idx / (nc * width);
    ray_map[b * nc * width * out_height * out_width +
            cw * out_height * out_width + invalid_idx] = static_cast<T>(0);
  }
  const int circle_map_idx = bs * dz;
  for (int32_t idx{0}; idx < circle_map_idx; ++idx) {
    int d = idx % dz;
    int b = idx / dz;
    circle_map[b * dz * out_height * out_width + d * out_height * out_width +
               invalid_idx] = static_cast<T>(0);
  }
  return 0;
}

template <typename T>
void gen_ring_ray_map(std::vector<double> const &Cam_Intrins,
                      std::vector<double> const &Cam_Extrins,
                      MatrixVTParameters const &matrix_vt_parameters,
                      std::vector<T> &circle_map_out,
                      std::vector<T> &ray_map_out) noexcept(false) {
  std::vector<float> const resize_scale = matrix_vt_parameters.resize_list;
  std::vector<CamCrop> const rect = matrix_vt_parameters.rect_list;
  std::vector<float> const voxel_size = matrix_vt_parameters.voxel_size;
  std::vector<float> const netRange = matrix_vt_parameters.net_range;
  int32_t batch_szie = matrix_vt_parameters.batch_size;
  int32_t sd_x = matrix_vt_parameters.sd_x;
  int32_t sd_y = matrix_vt_parameters.sd_y;
  int32_t n_cam = matrix_vt_parameters.n_cam;
  int32_t dz = matrix_vt_parameters.dz;
  int32_t width = matrix_vt_parameters.width;
  int32_t height = matrix_vt_parameters.height;
  int32_t out_width = matrix_vt_parameters.out_width;
  int32_t out_height = matrix_vt_parameters.out_height;
  int16_t scale = matrix_vt_parameters.scale;
  circle_map_out.resize(batch_szie * dz * out_height * out_width);
  ray_map_out.resize(batch_szie * n_cam * width * out_height * out_width);

  circle_map_out.assign(circle_map_out.size(), 0);
  ray_map_out.assign(ray_map_out.size(), 0);

  std::vector<double> augMat;
  std::vector<double> egoMat;
  augMat.resize(n_cam * 16);
  egoMat.resize(n_cam * 16);
  gen_aug_ego_pose(augMat, egoMat, Cam_Intrins, Cam_Extrins, resize_scale,
                   rect);

  std::vector<CustomCoorPoint> coorPoint;
  coorPoint.resize(batch_szie * n_cam * dz * height * width);
  cal_coor_cpu(coorPoint.data(), augMat.data(), egoMat.data(), voxel_size,
               netRange, sd_x, sd_y, n_cam, dz, width, height);
  auto invalid_x = static_cast<int32_t>(-netRange[0] / voxel_size[0]);
  auto invalid_y = static_cast<int32_t>(-netRange[1] / voxel_size[1]);
  std::vector<CustomCoorPoint> coorPoint_MeanH;
  std::vector<CustomCoorPoint> coorPoint_Permute;
  std::vector<int32_t> idx_h;
  coorPoint_MeanH.resize(batch_szie * n_cam * dz * width);
  coorPoint_Permute.resize(batch_szie * n_cam * dz * width);
  idx_h.resize(batch_szie * n_cam * dz * width);
  coor_reduce_hdim_cpu(coorPoint_MeanH.data(), coorPoint.data(), batch_szie,
                       n_cam, dz, height, width);
  coor_permute_cpu(coorPoint_Permute.data(), coorPoint_MeanH.data(),
                   idx_h.data(), batch_szie, n_cam, dz, width, out_height,
                   out_width, invalid_x, invalid_y);
  coor_map_cpu(circle_map_out.data(), ray_map_out.data(), idx_h.data(),
               batch_szie, dz, n_cam, width, out_height, out_width, scale,
               invalid_x, invalid_y);
}

}  // namespace perception
}  // namespace autodrive

namespace boost::json {

inline void tag_invoke(boost::json::value_from_tag, boost::json::value &jv,
                       autodrive::perception::CamCrop const &c) {
  jv = {{"x", c.x}, {"y", c.y}, {"x1", c.x1}, {"y1", c.y1}};
}

inline autodrive::perception::CamCrop tag_invoke(
    boost::json::value_to_tag<autodrive::perception::CamCrop>,
    boost::json::value const &jv) {
  auto const &obj = jv.as_object();
  return autodrive::perception::CamCrop{
      boost::json::value_to<float>(obj.at("x")),
      boost::json::value_to<float>(obj.at("y")),
      boost::json::value_to<float>(obj.at("x1")),
      boost::json::value_to<float>(obj.at("y1"))};
}

inline void tag_invoke(value_from_tag, value &jv,
                       autodrive::perception::MatrixVTParameters const &i) {
  jv = {{"batch_size", value_from(i.batch_size)},
        {"n_cam", value_from(i.n_cam)},
        {"dz", value_from(i.dz)},
        {"height", value_from(i.height)},
        {"width", value_from(i.width)},
        {"out_width", value_from(i.out_width)},
        {"out_height", value_from(i.out_height)},
        {"sd_x", value_from(i.sd_x)},
        {"sd_y", i.sd_y},
        {"scale", i.scale},
        {"voxel_size", value_from(i.voxel_size)},
        {"net_range", value_from(i.net_range)},
        {"rect_list", value_from(i.rect_list)},
        {"resize_list", value_from(i.resize_list)}};
}

inline autodrive::perception::MatrixVTParameters tag_invoke(
    boost::json::value_to_tag<autodrive::perception::MatrixVTParameters>,
    boost::json::value const &jv) {
  auto const &obj = jv.as_object();
  return autodrive::perception::MatrixVTParameters{
      boost::json::value_to<int32_t>(obj.at("batch_size")),
      boost::json::value_to<int32_t>(obj.at("n_cam")),
      boost::json::value_to<int32_t>(obj.at("dz")),
      boost::json::value_to<int32_t>(obj.at("height")),
      boost::json::value_to<int32_t>(obj.at("width")),
      boost::json::value_to<int32_t>(obj.at("out_width")),
      boost::json::value_to<int32_t>(obj.at("out_height")),
      boost::json::value_to<int32_t>(obj.at("sd_x")),
      boost::json::value_to<int32_t>(obj.at("sd_y")),
      boost::json::value_to<int16_t>(obj.at("scale")),
      boost::json::value_to<std::vector<float>>(obj.at("voxel_size")),
      boost::json::value_to<std::vector<float>>(obj.at("net_range")),
      boost::json::value_to<std::vector<autodrive::perception::CamCrop>>(
          obj.at("rect_list")),
      boost::json::value_to<std::vector<float>>(obj.at("resize_list"))};
}
}  // namespace boost::json

#endif  // PRO_MAT_INCLUDE_MATRIXVT_H_
